import { toPinyin } from "@/utils/pinyin";
import type { Key, MapTreeDataResult, TreeItem } from "../components/VirtualScroll";

export const selectTreeData = function <K extends Key>(treeData: MapTreeDataResult<K>, checked = true) {
  const [, data_map] = treeData;
  return (selected: K[] = [], key: K) => {
    const set = new Set(selected);
    const curr = data_map.get(key);
    if (checked) {
      set.add(key);
      data_map.forEach((item) => {
        if (item.parent_keys.includes(key)) {
          set.add(item.key);
        }
      });
      for (let i = curr?.parent_keys?.length - 1; i >= 0; i--) {
        const _key = curr.parent_keys[i];
        const element = data_map.get(_key);
        const select = !element.children.find((c) => !set.has(c.key));
        if (select) {
          set.add(_key);
        }
      }
    } else {
      set.delete(key);
      data_map.forEach((item) => {
        if (item.parent_keys.includes(key)) {
          set.delete(item.key);
        }
      });
      curr?.parent_keys?.forEach((k) => {
        set.delete(k);
      });
    }
    return [...set];
  };
};

export const pruneTreeData = function <K extends Key>(treeData: MapTreeDataResult<K>) {
  const [, data_map] = treeData;
  return (selected: K[]) => {
    const set = new Set(selected);
    return selected.filter((s) => {
      const item = data_map.get(s);
      return !item.parent_keys.find((k) => set.has(k));
    });
  };
};

export const searchTreeData = function <K extends Key>(
  keyword = "",
  treeData: MapTreeDataResult<K>,
  filter?: (item: TreeItem<K>) => boolean,
  render = (item: TreeItem<K>): any => item
): MapTreeDataResult<K> {
  keyword = keyword.trim();
  const visibles = new Map<K, TreeItem<K>>();
  const map = treeData[1];
  map.forEach((v, k) => {
    if (visibles.has(k)) {
      return;
    }

    if ((!keyword || v.label.includes(keyword)) && (!filter || filter(v))) {
      do {
        visibles.set(v.key, v);
        v = map.get(v.parent_key);
      } while (v);
    }
  });
  function loop(items: TreeItem<K>[]) {
    return items
      ?.map((item) => {
        if (visibles.has(item.key)) {
          return {
            ...render(item),
            children: loop(item.children),
          };
        }
      })
      .filter((a) => a);
  }
  return [loop(treeData[0]), visibles];
};

interface TreeItemMapKey<O> extends Partial<Record<keyof TreeItem, keyof O | string>> {
  __with_pinyin__?: boolean;
  /** 保留源对象中其余属性 */
  __with_rest__?: boolean;
}

export const mapTreeData = function <T extends Key, O>(
  options: O[] = [],
  {
    key = "key",
    value = "value",
    title = "title",
    label = "label",
    parent_key = "parent_key",
    children = "children",
    __with_pinyin__ = false,
    __with_rest__ = false,
  }: TreeItemMapKey<O>
): MapTreeDataResult<T, O> {
  let with_children = false;
  const children_map = new Map<T, (TreeItem<T> & O)[]>();
  const item_map = new Map<T, TreeItem<T> & O>();
  function loop(options: O[], pid?: T, pids: T[] = []) {
    const result: Array<TreeItem<T> & O> = [];
    for (let i = 0; i < options.length; i++) {
      const item = options[i];
      const {
        [key]: _key,
        [value]: _value,
        [title]: _title,
        [label]: _label,
        [parent_key]: _parent_key,
        [children]: _children,
        ...rest
      } = item as Record<string, string | number | O[]>;
      const obj = {
        ...(__with_rest__ ? rest : {}),
        key: _key,
        value: _value || _key,
        title: _title,
        label: _label || _title,
        parent_key: _parent_key,
        parent_keys: pids,
        children: _children ? loop(_children as O[], _key as T, pids.concat(_key as T)) : undefined,
      } as unknown as TreeItem<T> & O;
      if (__with_pinyin__) {
        obj.label = `${obj.label} ${toPinyin(obj.label)}`;
      }
      if (typeof pid != "undefined") {
        // 如果有pid说明存在children，就不在进行整理
        with_children = true;
        obj.parent_key = pid;
      } else if (typeof obj.parent_key != "undefined") {
        // 如果有parent_key 需要整理一下children数组
        const _children_ = children_map.get(obj.parent_key) || [];
        const parent = item_map.get(obj.parent_key);
        children_map.set(
          obj.parent_key,
          _children_.concat(obj).map((c) => {
            c.parent_keys = (parent?.parent_keys || []).concat(parent ? [obj.parent_key] : []);
            return c;
          })
        );
      } else {
        // 顶层不处理
      }
      item_map.set(obj.key, obj);
      result.push(obj);
    }
    return result.length > 0 ? result : undefined;
  }
  let tree = loop(options);
  if (!with_children) {
    tree = tree?.filter((o) => {
      o.children = children_map.get(o.key);
      // 严格检查，避免服务端返回部分数据导致用parentid检测根节点错误
      return typeof o.parent_key === "undefined" || !item_map.has(o.parent_key);
    });
  }
  return [tree || [], item_map] as const;
};
